use std::thread;
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::Mutex;

enum Message {
    NewJob(Job), // 定义一个枚举类型，表示线程池可以接收的消息类型，包括新任务和终止信号。
    Terminate,   // 当收到 Terminate 消息时，工作者线程将结束运行。
}

pub struct ThreadPool { // 定义 ThreadPool 结构体，代表线程池。
    workers: Vec<Worker>, // 存储工作者线程的向量。
    sender: mpsc::Sender<Message>, // 用于发送消息到工作者线程的发送端。
}

impl ThreadPool {
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);
        let (sender, receiver) = mpsc::channel(); // 创建一个通道用于线程间通信。
        let receiver = Arc::new(Mutex::new(receiver)); // 使用 Arc 和 Mutex 包装 receiver，以便在多线程中安全使用。

        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool { workers, sender }
    }
    
    pub fn execute<F>(&self, f: F)
    where 
        F: FnOnce() + Send + 'static, // 约束 F 为 FnOnce，即只能被调用一次的闭包，并且可以跨线程发送和 'static 生命周期。
    {
        let job = Box::new(f); // 将闭包 f 转换为 Box，满足 Job 类型。
        self.sender.send(Message::NewJob(job)).unwrap(); // 发送新任务消息到通道。
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        println!("Sending Terminate message to all workers.");
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).unwrap();
        }
        println!("Shutting down all workers.");

        for worker in &mut self.workers {
            println!("Shutting down worker {}", worker.id);
            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap(); // 等待线程结束。
            }
        }
    }
}

struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

// 定义 FnBox 特性，允许将闭包转换为 Box。
trait FnBox {
    fn call_box(self: Box<Self>); // 将闭包调用并执行。
}

// 为所有 FnOnce 闭包实现 FnBox 特性。
impl <F: FnOnce()> FnBox for F {
    fn call_box(self: Box<F>) {
        (*self)() // 调用闭包。
    }
}

// 定义 Job 类型，它是实现了 FnBox 特性的 trait 对象，可以发送到其他线程。
type Job = Box<dyn FnBox + Send + 'static>;

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
        let thread = thread::spawn(move || loop { // 创建一个新的线程，并移动 receiver 到线程中。
            let message = receiver.lock().unwrap().recv().unwrap(); // 锁定 Mutex 并接收消息。

            match message { // 匹配接收到的消息。
                Message::NewJob(job) => { // 如果是新任务消息。
                    println!("Worker {} got a job; executing.", id);
                    job.call_box(); // 执行任务。
                },
                Message::Terminate => {
                    println!("Worker {} was terminated.", id);
                    break;
                }
            }
        });

        Worker { id, thread: Some(thread) }
    }
}